/*
 * Copyright (C) 2016 The Android Open Source Project
 *
 * Permission is hereby granted, free of charge, to any person
 * obtaining a copy of this software and associated documentation
 * files (the "Software"), to deal in the Software without
 * restriction, including without limitation the rights to use, copy,
 * modify, merge, publish, distribute, sublicense, and/or sell copies
 * of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be
 * included in all copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
 * EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
 * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
 * NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
 * BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
 * ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
#include <common.h>
#include <image.h>
#include <android_image.h>
#include <malloc.h>
#include <memalign.h>
#include <mapmem.h>
#include <errno.h>
#include <command.h>
#include <mmc.h>
#include <blk.h>
#include <part.h>
#include <android_avb/libavb_ab.h>
#include <config.h>
#include <command.h>
#include <environment.h>
#include <linux/types.h>
#include <api_public.h>
#include <android_avb/libavb.h>
#include <android_avb/avb_util.h>
#include <android_avb/avb_ab_flow.h>


bool avb_ab_data_verify_and_byteswap(const AvbABData* src, AvbABData* dest)
{
	/* Ensure magic is correct. */
	if (avb_safe_memcmp(src->magic, AVB_AB_MAGIC, AVB_AB_MAGIC_LEN) != 0) {
		avb_error("Magic is incorrect.\n");
		return false;
	}

	avb_memcpy(dest, src, sizeof(AvbABData));
	dest->crc32 = avb_be32toh(dest->crc32);

	/* Ensure we don't attempt to access any fields if the major version
	* is not supported.
	*/
	if (dest->version_major > AVB_AB_MAJOR_VERSION) {
		avb_error("No support for given major version.\n");
		return false;
	}

	/* Bail if CRC32 doesn't match. */
	if (dest->crc32 !=
		avb_crc32((const uint8_t*)dest, sizeof(AvbABData) - sizeof(uint32_t) - 1)) {
		avb_error("CRC32 does not match.\n");
		return false;
	}

	return true;
}

void avb_ab_data_update_crc_and_byteswap(const AvbABData* src,
                                         AvbABData* dest)
{
	avb_memcpy(dest, src, sizeof(AvbABData));
	dest->crc32 = avb_htobe32(
		avb_crc32((const uint8_t*)dest, sizeof(AvbABData) - sizeof(uint32_t) - 1));
}

void avb_ab_data_init(AvbABData *data)
{
	avb_memset(data, '\0', sizeof(AvbABData));
	avb_memcpy(data->magic, AVB_AB_MAGIC, AVB_AB_MAGIC_LEN);
	data->nb_slot = 2;
	data->slot_suffix[0] = 'a';
	data->slot_suffix[1] = 'b';
	data->slot_suffix[2] = '\0';
	data->slot_suffix[3] = '\0';
	data->version_major = AVB_AB_MAJOR_VERSION;
	data->version_minor = AVB_AB_MINOR_VERSION;
	data->slots[0].priority = AVB_AB_MAX_PRIORITY;
	data->slots[0].tries_remaining = AVB_AB_MAX_TRIES_REMAINING;
	data->slots[0].successful_boot = 0;
	data->slots[1].priority = AVB_AB_MAX_PRIORITY - 1;
	data->slots[1].tries_remaining = AVB_AB_MAX_TRIES_REMAINING;
	data->slots[1].successful_boot = 0;
}

/* The AvbABData struct is stored 2048 bytes into the 'misc' partition
 * following the 'struct bootloader_message' field. The struct is
 * compatible with the guidelines in bootable/recovery/bootloader.h -
 * e.g. it is stored in the |slot_suffix| field, starts with a
 * NUL-byte, and is 32 bytes long.
 */
#define AB_METADATA_MISC_PARTITION_OFFSET 2048

AvbIOResult avb_ab_data_read(AvbABOps* ab_ops, AvbABData* data)
{
	AvbOps* ops = ab_ops->ops;
	ALLOC_CACHE_ALIGN_BUFFER(u8, serialized, 512);
	size_t num_bytes_read;

	memset(serialized,0,sizeof(u8) * 512);
	ops->read_from_partition(ops,
				"misc",
				AB_METADATA_MISC_PARTITION_OFFSET,
				512,
				(AvbABData*)serialized,
				&num_bytes_read);
	if (!avb_ab_data_verify_and_byteswap((AvbABData*)serialized, data)) {
	avb_error(
		"Error validating A/B metadata from disk. "
		"Resetting and writing new A/B metadata to disk.\n");
		avb_ab_data_init(data);
		return avb_ab_data_write(ab_ops, data);
	}
	return AVB_IO_RESULT_OK;
}

AvbIOResult avb_ab_data_write(AvbABOps* ab_ops, const AvbABData* data) {
	AvbOps* ops = ab_ops->ops;
	AvbIOResult io_ret;
	ALLOC_CACHE_ALIGN_BUFFER(u8, temp, 512);
	memset(temp,0,sizeof(u8) * 512);

	avb_ab_data_update_crc_and_byteswap(data, (AvbABData*)temp);
	io_ret = ops->write_to_partition(ops,
                                   "misc",
                                   AB_METADATA_MISC_PARTITION_OFFSET,
                                   512,
                                   temp);
	if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
		return AVB_IO_RESULT_ERROR_OOM;
	} else if (io_ret != AVB_IO_RESULT_OK) {
		avb_error("Error writing A/B metadata.\n");
		return AVB_IO_RESULT_ERROR_IO;
	}
	return AVB_IO_RESULT_OK;
}

static bool slot_is_bootable(AvbABSlotData* slot)
{
	return slot->priority > 0 &&
		(slot->successful_boot || (slot->tries_remaining > 0));
}

static void slot_set_unbootable(AvbABSlotData* slot)
{
	slot->priority = 0;
	slot->tries_remaining = 0;
	slot->successful_boot = 0;
}

/* Ensure all unbootable and/or illegal states are marked as the
 * canonical 'unbootable' state, e.g. priority=0, tries_remaining=0,
 * and successful_boot=0.
 */
static void slot_normalize(AvbABSlotData* slot)
{
	if (slot->priority > 0) {
		if (slot->tries_remaining == 0 && !slot->successful_boot) {
		/* We've exhausted all tries -> unbootable. */
		slot_set_unbootable(slot);
	}
	if (slot->tries_remaining > 0 && slot->successful_boot) {
		/* Illegal state - avb_ab_mark_slot_successful() will clear
		 * tries_remaining when setting successful_boot.
		 */
		slot_set_unbootable(slot);
	}
	} else {
		slot_set_unbootable(slot);
	}
}

static const char* slot_suffixes[2] = {"_a", "_b"};

/* Helper function to load metadata - returns AVB_IO_RESULT_OK on
 * success, error code otherwise.
 */
AvbIOResult load_metadata(AvbABOps *ab_ops,
				AvbABData *ab_data,
				AvbABData *ab_data_orig)
{
	AvbIOResult io_ret;
	io_ret = ab_ops->read_ab_metadata(ab_ops,ab_data);
	if(io_ret != AVB_IO_RESULT_OK) {
		avb_error("I/O error while loading A/B metadata.\n");
		return io_ret;
	}

	*ab_data_orig = *ab_data;
	slot_normalize(&ab_data->slots[0]);
	slot_normalize(&ab_data->slots[1]);
	return AVB_IO_RESULT_OK;
}


/* Writes A/B metadata to disk only if it has changed - returns
 * AVB_IO_RESULT_OK on success, error code otherwise.
 */
static AvbIOResult save_metadata_if_changed(AvbABOps* ab_ops,
                                            AvbABData* ab_data,
                                            AvbABData* ab_data_orig)
{
	if (avb_safe_memcmp(ab_data, ab_data_orig, sizeof(AvbABData)) != 0) {
		avb_debug("Writing A/B metadata to disk.\n");
		return ab_ops->write_ab_metadata(ab_ops, ab_data);
	}
	return AVB_IO_RESULT_OK;
}

AvbABFlowResult avb_ab_flow(AvbABOps* ab_ops,
                            const char* const* requested_partitions,
                            AvbSlotVerifyFlags flags,
                            AvbHashtreeErrorMode hashtree_error_mode,
                            AvbSlotVerifyData** out_data)
{
	AvbOps* ops = ab_ops->ops;
	AvbSlotVerifyData* slot_data[2] = {NULL, NULL};
	AvbSlotVerifyData* data = NULL;
	AvbABFlowResult ret;
	AvbABData ab_data, ab_data_orig;
	size_t slot_index_to_boot, n;
	AvbIOResult io_ret;
	bool saw_and_allowed_verification_error = false;
	io_ret = load_metadata(ab_ops, &ab_data, &ab_data_orig);
	if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
		ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
		goto out;
	} else if (io_ret != AVB_IO_RESULT_OK) {
		ret = AVB_AB_FLOW_RESULT_ERROR_IO;
		goto out;
	}

	/* Validate all bootable slots. */
	for (n = 0; n < 2; n++) {
		if (slot_is_bootable(&ab_data.slots[n])) {
			AvbSlotVerifyResult verify_result;
			bool set_slot_unbootable = false;

			verify_result = avb_slot_verify(ops,
                                      requested_partitions,
                                      slot_suffixes[n],
                                      flags,
                                      hashtree_error_mode,
                                      &slot_data[n]);
			switch (verify_result) {
			case AVB_SLOT_VERIFY_RESULT_ERROR_OOM:
				ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
				goto out;

			case AVB_SLOT_VERIFY_RESULT_ERROR_IO:
				ret = AVB_AB_FLOW_RESULT_ERROR_IO;
				goto out;

			case AVB_SLOT_VERIFY_RESULT_OK:
				break;

			case AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_METADATA:
			case AVB_SLOT_VERIFY_RESULT_ERROR_UNSUPPORTED_VERSION:
			/* Even with AVB_SLOT_VERIFY_FLAGS_ALLOW_VERIFICATION_ERROR
		 	 * these mean game over.
		 	 */
				set_slot_unbootable = true;
				break;

			/* explicit fallthrough. */
			case AVB_SLOT_VERIFY_RESULT_ERROR_VERIFICATION:
			case AVB_SLOT_VERIFY_RESULT_ERROR_ROLLBACK_INDEX:
			case AVB_SLOT_VERIFY_RESULT_ERROR_PUBLIC_KEY_REJECTED:
				if (flags & AVB_SLOT_VERIFY_FLAGS_ALLOW_VERIFICATION_ERROR) {
				/* Do nothing since we allow this. */
					avb_debugv("Allowing slot ",
					slot_suffixes[n],
					" which verified "
					"with result ",
					avb_slot_verify_result_to_string(verify_result),
					" because "
					"AVB_SLOT_VERIFY_FLAGS_ALLOW_VERIFICATION_ERROR "
					"is set.\n",
					NULL);
					saw_and_allowed_verification_error = true;
				} else {
					set_slot_unbootable = true;
				}
				break;

			case AVB_SLOT_VERIFY_RESULT_ERROR_INVALID_ARGUMENT:
				ret = AVB_AB_FLOW_RESULT_ERROR_INVALID_ARGUMENT;
				goto out;
			/* Do not add a 'default:' case here because of -Wswitch. */
			}

			if (set_slot_unbootable) {
				avb_errorv("Error verifying slot ",
					slot_suffixes[n],
					" with result ",
					avb_slot_verify_result_to_string(verify_result),
					" - setting unbootable.\n",
					NULL);
				slot_set_unbootable(&ab_data.slots[n]);
			}
		}
	}

	if (slot_is_bootable(&ab_data.slots[0]) && slot_is_bootable(&ab_data.slots[1])) {
		if (ab_data.slots[1].priority > ab_data.slots[0].priority) {
			slot_index_to_boot = 1;
		} else {
			slot_index_to_boot = 0;
		}
		} else if (slot_is_bootable(&ab_data.slots[0])) {
			slot_index_to_boot = 0;
		} else if (slot_is_bootable(&ab_data.slots[1])) {
			slot_index_to_boot = 1;
		} else {
		/* No bootable slots! */
			avb_error("No bootable slots found.\n");
			ret = AVB_AB_FLOW_RESULT_ERROR_NO_BOOTABLE_SLOTS;
			goto out;
		}

		/* Update stored rollback index such that the stored rollback index
		 * is the largest value supporting all currently bootable slots. Do
		 * this for every rollback index location.
		 */
		for (n = 0; n < AVB_MAX_NUMBER_OF_ROLLBACK_INDEX_LOCATIONS; n++) {
			uint64_t rollback_index_value = 0;

			if (slot_data[0] != NULL && slot_data[1] != NULL) {
				uint64_t a_rollback_index = slot_data[0]->rollback_indexes[n];
				uint64_t b_rollback_index = slot_data[1]->rollback_indexes[n];
				rollback_index_value =
				(a_rollback_index < b_rollback_index ? a_rollback_index
								: b_rollback_index);
			} else if (slot_data[0] != NULL) {
				rollback_index_value = slot_data[0]->rollback_indexes[n];
			} else if (slot_data[1] != NULL) {
				rollback_index_value = slot_data[1]->rollback_indexes[n];
			}

			if (rollback_index_value != 0) {
				uint64_t current_rollback_index_value;
				io_ret = ops->read_rollback_index(ops, n, &current_rollback_index_value);
				if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
					ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
					goto out;
			} else if (io_ret != AVB_IO_RESULT_OK) {
				avb_error("Error getting rollback index for slot.\n");
				ret = AVB_AB_FLOW_RESULT_ERROR_IO;
				goto out;
			}
			if (current_rollback_index_value != rollback_index_value) {
				io_ret = ops->write_rollback_index(ops, n, rollback_index_value);
				if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
 					ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
					goto out;
				} else if (io_ret != AVB_IO_RESULT_OK) {
					avb_error("Error setting stored rollback index.\n");
          				ret = AVB_AB_FLOW_RESULT_ERROR_IO;
          				goto out;
				}
			}
		}
	}

	/* Finally, select this slot. */
	avb_assert(slot_data[slot_index_to_boot] != NULL);
	data = slot_data[slot_index_to_boot];
	slot_data[slot_index_to_boot] = NULL;
	if (saw_and_allowed_verification_error) {
		avb_assert(flags & AVB_SLOT_VERIFY_FLAGS_ALLOW_VERIFICATION_ERROR);
		ret = AVB_AB_FLOW_RESULT_OK_WITH_VERIFICATION_ERROR;
	} else {
		ret = AVB_AB_FLOW_RESULT_OK;
	}

	/* ... and decrement tries remaining, if applicable. */
	if (!ab_data.slots[slot_index_to_boot].successful_boot &&
		ab_data.slots[slot_index_to_boot].tries_remaining > 0) {
		ab_data.slots[slot_index_to_boot].tries_remaining -= 1;
	}

out:
	io_ret = save_metadata_if_changed(ab_ops, &ab_data, &ab_data_orig);
	if (io_ret != AVB_IO_RESULT_OK) {
		if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
		ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
	} else {
			ret = AVB_AB_FLOW_RESULT_ERROR_IO;
		}
		if (data != NULL) {
			avb_slot_verify_data_free(data);
			data = NULL;
		}
	}

	for (n = 0; n < 2; n++) {
		if (slot_data[n] != NULL) {
			avb_slot_verify_data_free(slot_data[n]);
		}
	}

	if (out_data != NULL) {
		*out_data = data;
	} else {
		if (data != NULL) {
			avb_slot_verify_data_free(data);
		}
	}

	return ret;
}

AvbABFlowResult avb_ab_slot_select(AvbABOps* ab_ops,char select_slot[])
{
	AvbABFlowResult ret = 0;
	AvbIOResult io_ret = 0;
	AvbABData ab_data, ab_data_orig;
	size_t slot_index_to_boot;
	io_ret = load_metadata(ab_ops, &ab_data, &ab_data_orig);
	if (io_ret == AVB_IO_RESULT_ERROR_OOM) {
		ret = AVB_AB_FLOW_RESULT_ERROR_OOM;
		goto out;
	} else if (io_ret != AVB_IO_RESULT_OK) {
		ret = AVB_AB_FLOW_RESULT_ERROR_IO;
		goto out;
	}
	if (slot_is_bootable(&ab_data.slots[0]) && slot_is_bootable(&ab_data.slots[1])) {
		if (ab_data.slots[1].priority > ab_data.slots[0].priority) {
			slot_index_to_boot = 1;
		} else {
			slot_index_to_boot = 0;
		}
	} else if(slot_is_bootable(&ab_data.slots[0])) {
		slot_index_to_boot = 0;
	} else if(slot_is_bootable(&ab_data.slots[1])) {
		slot_index_to_boot = 1;
	} else {
		avb_error("No bootable slots found.\n");
		ret = AVB_AB_FLOW_RESULT_ERROR_NO_BOOTABLE_SLOTS;
		goto out;
	}
	if (slot_index_to_boot == 0) {
		select_slot[0] = '_';
		select_slot[1] = 'a';
		select_slot[2] = '\0';
	} else if(slot_index_to_boot == 1) {
		select_slot[0] = '_';
		select_slot[1] = 'b';
		select_slot[2] = '\0';
	}
out:
	return ret;
}

AvbIOResult avb_ab_mark_slot_active(AvbABOps *ab_ops,
                                    unsigned int slot_number)
{
	AvbABData ab_data, ab_data_orig;
	unsigned int other_slot_number;
	AvbIOResult ret;

	if( slot_number > 1) {
		avb_error("slot_number can not > 1");
		return -1;
	}
	ret = load_metadata(ab_ops, &ab_data, &ab_data_orig);
	if (ret != AVB_IO_RESULT_OK) {
		goto out;
	}
	/* Make requested slot top priority, unsuccessful, and with max tries. */
	ab_data.slots[slot_number].priority = AVB_AB_MAX_PRIORITY;
	ab_data.slots[slot_number].tries_remaining = AVB_AB_MAX_TRIES_REMAINING;
	ab_data.slots[slot_number].successful_boot = 0;

	/* Ensure other slot doesn't have as high a priority. */
	other_slot_number = 1 - slot_number;
	if (ab_data.slots[other_slot_number].priority == AVB_AB_MAX_PRIORITY) {
		ab_data.slots[other_slot_number].priority = AVB_AB_MAX_PRIORITY - 1;
	}

	ret = AVB_IO_RESULT_OK;

out:
	if (ret == AVB_IO_RESULT_OK) {
		ret = save_metadata_if_changed(ab_ops, &ab_data, &ab_data_orig);
	}
	return ret;
}

AvbIOResult avb_ab_mark_slot_unbootable(AvbABOps *ab_ops,
                                        unsigned int slot_number)
{
	AvbABData ab_data, ab_data_orig;
	AvbIOResult ret;

	if( slot_number > 1) {
		avb_error("slot_number can not > 1");
		return -1;
	}

	ret = load_metadata(ab_ops, &ab_data, &ab_data_orig);
	if (ret != AVB_IO_RESULT_OK) {
		goto out;
	}

	slot_set_unbootable(&ab_data.slots[slot_number]);

	ret = AVB_IO_RESULT_OK;

out:
	if (ret == AVB_IO_RESULT_OK) {
		ret = save_metadata_if_changed(ab_ops, &ab_data, &ab_data_orig);
	}
	return ret;
}

AvbIOResult avb_ab_mark_slot_successful(AvbABOps *ab_ops,
                                        unsigned int slot_number)
{
	AvbABData ab_data, ab_data_orig;
	AvbIOResult ret;

	if( slot_number > 1) {
		avb_error("slot_number can not > 1");
		return -1;
	}

	ret = load_metadata(ab_ops, &ab_data, &ab_data_orig);
	if (ret != AVB_IO_RESULT_OK) {
		goto out;
	}

	if (!slot_is_bootable(&ab_data.slots[slot_number])) {
		avb_error("Cannot mark unbootable slot as successful.\n");
		ret = AVB_IO_RESULT_OK;
		goto out;
	}

	ab_data.slots[slot_number].tries_remaining = 0;
	ab_data.slots[slot_number].successful_boot = 1;

	ret = AVB_IO_RESULT_OK;

out:
	if (ret == AVB_IO_RESULT_OK) {
		ret = save_metadata_if_changed(ab_ops, &ab_data, &ab_data_orig);
	}
	return ret;
}

const char *avb_ab_flow_result_to_string(AvbABFlowResult result)
{
	const char* ret = NULL;

	switch (result) {
	case AVB_AB_FLOW_RESULT_OK:
		ret = "OK";
		break;

	case AVB_AB_FLOW_RESULT_OK_WITH_VERIFICATION_ERROR:
		ret = "OK_WITH_VERIFICATION_ERROR";
		break;

	case AVB_AB_FLOW_RESULT_ERROR_OOM:
		ret = "ERROR_OOM";
		break;

	case AVB_AB_FLOW_RESULT_ERROR_IO:
		ret = "ERROR_IO";
		break;

	case AVB_AB_FLOW_RESULT_ERROR_NO_BOOTABLE_SLOTS:
		ret = "ERROR_NO_BOOTABLE_SLOTS";
		break;
	case AVB_AB_FLOW_RESULT_ERROR_INVALID_ARGUMENT:
		ret = "ERROR_INVALID_ARGUMENT";
		break;
	/* Do not add a 'default:' case here because of -Wswitch. */
	}

	if (ret == NULL) {
		avb_error("Unknown AvbABFlowResult value.\n");
		ret = "(unknown)";
	}

	return ret;
}
